'''
@author: Zhouhong Gu
@date: 2021/12/29
@idea:  拟合SEIR model
'''
import numpy as np
import pandas as pd
from scipy.integrate import solve_ivp
from scipy.optimize import minimize
import matplotlib.pyplot as plt
import os
from config import PathVar
from datetime import timedelta, datetime


class Learner(object):
    def __init__(self, loss, city=None, path='./data/02_SZ_DailyCases.csv', parameter_output= None):
        self.loss = loss
        self.city = city
        self.path = path
        self.parameter_output = parameter_output
        self.pathvar = PathVar()

    def input_range_parameters(self, range_parameters: dict):
        # 超参数
        self.E_ratio = range_parameters['E_ratio']
        self.Ehq_ratio = range_parameters['Ehq_ratio']
        self.Ecq_ratio = range_parameters['Ecq_ratio']

        self.A_loss_rate = range_parameters['A_loss_rate']
        self.I_loss_rate = range_parameters['I_loss_rate']
        self.R_loss_rate = 1 - self.A_loss_rate - self.I_loss_rate

        self.t_1 = range_parameters['t_1']
        self.t_2 = range_parameters['t_2']
        self.q_1_min = range_parameters['q_1_min']
        self.q_1_max = range_parameters['q_1_max']
        self.q_2_min = range_parameters['q_2_min']
        self.q_2_max = range_parameters['q_2_max']
        self.q_3_min = range_parameters['q_3_min']
        self.q_3_max = range_parameters['q_3_max']

        self.p_1_min = range_parameters['p_1_min']
        self.p_1_max = range_parameters['p_1_max']
        self.p_2_min = range_parameters['p_2_min']
        self.p_2_max = range_parameters['p_2_max']
        self.p_3_min = range_parameters['p_3_min']
        self.p_3_max = range_parameters['p_3_max']

        self.beta_1_min = range_parameters['beta_1_min']
        self.beta_1_max = range_parameters['beta_1_max']
        self.beta_2_min = range_parameters['beta_2_min']
        self.beta_2_max = range_parameters['beta_2_max']
        self.beta_3_min = range_parameters['beta_3_min']
        self.beta_3_max = range_parameters['beta_3_max']

        self.Lambda_min = range_parameters['Lambda_min']
        self.Lambda_max = range_parameters['Lambda_max']
        self.eposilon_min = range_parameters['eposilon_min']
        self.eposilon_max = range_parameters['eposilon_max']
        self.theta_min = range_parameters['theta_min']
        self.theta_max = range_parameters['theta_max']
        self.alpha_min = range_parameters['alpha_min']
        self.alpha_max = range_parameters['alpha_max']
        self.eta_min = range_parameters['eta_min']
        self.eta_max = range_parameters['eta_min']
        self.gamma_min = range_parameters['gamma_min']
        self.gamma_max = range_parameters['gamma_max']
        self.Q_min = range_parameters['Q_min']
        self.Q_max = range_parameters['Q_max']

        # 初始值
        self.predict_range = range_parameters['predict_range']

    def input_start_parameters(self, start_parameters: dict):
        self.q_1_0 = start_parameters['q_1']  # 隔离率_1
        self.t_1 = start_parameters['t_1']  # q_1持续时间
        self.q_2_0 = start_parameters['q_2']  # 隔离率_2
        self.t_2 = start_parameters['t_2']  # q_2持续时间
        self.q_3_0 = start_parameters['q_3']  # 隔离率_3
        self.p_1_0 = start_parameters['p_1']  # 管控率
        self.p_2_0 = start_parameters['p_2']  # 管控率
        self.p_3_0 = start_parameters['p_3']  # 管控率
        self.Lambda_0 = start_parameters['Lambda']  # 隔离解除率
        self.beta_1_0 = start_parameters['beta_1']  # 感染者接触数——后面考虑建一个一次方程
        self.beta_2_0 = start_parameters['beta_2']  # 感染者接触数——后面考虑建一个一次方程
        self.beta_3_0 = start_parameters['beta_3']  # 感染者接触数——后面考虑建一个一次方程
        self.eposilon_0 = start_parameters['eposilon']  # 潜伏期传染力系数
        self.theta_0 = start_parameters['theta']  # 无症状感染者传染力系数
        self.alpha_0 = start_parameters['alpha']  # 隔离率
        self.eta_0 = start_parameters['eta']  # 感染者中确诊患者比例
        self.gamma_0 = start_parameters['gamma']  # 感染者移出率
        self.Q_0 = start_parameters['Q']  # 每天隔离人数

    def input_start_variables(self, start_variables: dict):

        self.s_0 = start_variables['s']
        self.sq_0 = start_variables['sq']
        self.e_0 = start_variables['e']
        self.ehq_0 = start_variables['ehq']
        self.ecq_0 = start_variables['ecq']
        self.i_0 = start_variables['i']
        self.a_0 = start_variables['a']
        self.r_0 = start_variables['r']
        self.N = sum([self.s_0, self.sq_0, self.e_0, self.ehq_0, self.ecq_0, self.i_0, self.a_0, self.r_0])

    def load_I(self):
        '''
        获取累计确诊
        '''
        df = pd.read_csv(self.path)
        if self.city:
            df = df[df['City'] == self.city]
        df.set_index(["Date"], inplace=True)
        dff = df["I"]
        return dff.T

    def load_A(self):
        '''
        获取累计确诊
        '''
        df = pd.read_csv(self.path)
        if self.city:
            df = df[df['City'] == self.city]
        df.set_index(["Date"], inplace=True)
        dff = df["A"]
        return dff.T

    def load_E(self, ratio):
        '''
        估计潜在患者
        '''
        df = pd.read_csv(self.path)
        if self.city:
            df = df[df['City'] == self.city]
        df.set_index(["Date"], inplace=True)
        dff = df["A"] + df['I']
        dfff = dff * ratio
        return dfff.T

    def load_Ehq(self, ratio):
        '''
        估计潜在患者
        '''
        df = pd.read_csv(self.path)
        if self.city:
            df = df[df['City'] == self.city]
        df.set_index(["Date"], inplace=True)
        dff = df["A"] + df['I']
        dfff = dff * ratio
        return dfff.T

    def load_Ecq(self, ratio):
        '''
        估计潜在患者
        '''
        df = pd.read_csv(self.path)
        if self.city:
            df = df[df['City'] == self.city]
        df.set_index(["Date"], inplace=True)
        dff = df["A"] + df['I']
        dfff = dff * ratio
        return dfff.T

    def load_Sq(self):
        df = pd.read_csv(self.path)
        if self.city:
            df = df[df['City'] == self.city]
        df.set_index(["Date"], inplace=True)
        index = df.index
        length = len(df)
        return pd.Series([self.Q_0 * (index + 1) for index in range(length)], index=index)

    def load_R(self):
        '''
        获取累计治愈
        '''
        df = pd.read_csv(self.path)
        if self.city:
            df = df[df['City'] == self.city]
        df.set_index(["Date"], inplace=True)
        dff = df['R']
        return dff.T

    def extend_index(self, index, new_size):
        values = index.values
        current = datetime.strptime(index[-1], '%Y/%m/%d')
        while len(values) < new_size:
            current = current + timedelta(days=1)
            values = np.append(values, datetime.strftime(current, '%Y/%m/%d'))
        return values

    def predict(self, q_1, q_2, q_3, p_1,p_2,p_3, beta_1,beta_2,beta_3, eposilon, theta, alpha, eta, gamma, Q, Lambda, S, Sq, E, Ehq, Ecq, A, I,
                R, s_0, sq_0, e_0, ehq_0, ecq_0, i_0, a_0, r_0):
        new_index = self.extend_index(S.index, self.predict_range)
        size = len(new_index)

        def SEIR(t, y):
            '''
            return: ds,de,di,dr
            '''
            S = y[0]
            Sq = y[1]
            E = y[2]
            Ehq = y[3]
            Ecq = y[4]
            I = y[5]
            A = y[6]
            R = y[7]
            q = q_1 if t < self.t_1 else q_2 if t < self.t_2 + self.t_1 else q_3
            p = p_1 if t < self.t_1 else p_2 if t< self.t_2 + self.t_1 else p_3
            beta = beta_1 if t<self.t_1 else beta_2 if t<self.t_2 + self.t_1 else beta_3
            Ehq_in = (1 - q) * p * beta * (I + eposilon * E + theta * A) * S
            Ehq_out = alpha * eta * Ehq + alpha * (1 - eta) * Ehq

            E_in = (1 - q) * (1 - p) * beta * (I + eposilon * E + theta * A) * S
            E_out = alpha * eta * E + alpha * (1 - eta) * E

            Ecq_in = q * beta * (I + eposilon * E + theta * A) * S
            Ecq_out = alpha * eta * Ecq + alpha * (1 - eta) * Ecq

            S_in = Lambda * Sq
            S_out = Ehq_in + E_in + Ecq_in + Q

            Sq_out = Lambda * Sq
            Sq_in = Q

            I_in = alpha * eta * (Ehq + E + Ecq)
            A_in = alpha * (1 - eta) * (Ehq + E + Ecq)

            I_out = gamma * I
            A_out = gamma * A

            R_in = I_out + A_out

            return [S_in - S_out,
                    Sq_in - Sq_out,
                    E_in - E_out,
                    Ehq_in - Ehq_out,
                    Ecq_in - Ecq_out,
                    I_in - I_out,
                    A_in - A_out,
                    R_in]

        extended_S = np.concatenate((S.values, [None] * (size - len(S.values))))
        extended_Sq = np.concatenate((Sq.values, [None] * (size - len(Sq.values))))
        extended_E = np.concatenate((E.values, [None] * (size - len(E.values))))
        extended_Ehq = np.concatenate((Ehq.values, [None] * (size - len(Ehq.values))))
        extended_Ecq = np.concatenate((Ecq.values, [None] * (size - len(Ecq.values))))
        extended_I = np.concatenate((I.values, [None] * (size - len(I.values))))
        extended_A = np.concatenate((A.values, [None] * (size - len(A.values))))
        extended_R = np.concatenate((R.values, [None] * (size - len(R.values))))
        return new_index, extended_S, extended_Sq, extended_E, extended_Ehq, extended_Ecq, extended_A, extended_I, extended_R, solve_ivp(
            SEIR, [0, size], [s_0, sq_0, e_0, ehq_0, ecq_0, i_0, a_0, r_0], t_eval=np.arange(0, size, 1))

    def train(self):
        print('start training')
        R = self.load_R()
        E = self.load_E(self.E_ratio)
        Ehq = self.load_Ehq(self.Ehq_ratio)
        Ecq = self.load_Ecq(self.Ecq_ratio)
        A = self.load_A()
        I = self.load_I()
        Sq = self.load_Sq()
        S = self.s_0 - (I + A + E + Ecq + Ehq + R + Sq)  # 易感人数
        print('parameters loaded')
        print(f'学习率：I-{self.I_loss_rate}, A-{self.A_loss_rate}, R-{self.R_loss_rate}')
        # given 的数据只有A，I，R是确定的
        # q_1, q_2, q_3, p, beta, eposilon, theta, alpha, eta, gamma, Q, Lambda = point
        methods = ['L-BFGS-B', 'BFGS', 'Newton-CG', 'CG', 'Powell', 'Nelder-Mead']
        optimal = minimize(self.loss,
                           np.array([self.q_1_0, self.q_2_0, self.q_3_0, self.p_1_0, self.p_2_0, self.p_3_0,
                                     self.beta_1_0, self.beta_2_0, self.beta_3_0, self.eposilon_0,
                                     self.theta_0, self.alpha_0, self.eta_0, self.gamma_0, self.Q_0,
                                     self.Lambda_0]),
                           args=(
                               self.t_1, self.t_2, self.I_loss_rate, self.A_loss_rate, self.R_loss_rate, A, I, R,
                               int(self.s_0), int(self.sq_0), int(self.e_0), int(self.ehq_0), int(self.ecq_0),
                               int(self.i_0), int(self.a_0), int(self.r_0)),
                           method=methods[0],
                           bounds=[(self.q_1_min, self.q_1_max), (self.q_2_min, self.q_2_max),
                                   (self.q_3_min, self.q_3_max), (self.p_1_min, self.p_1_max),
                                   (self.p_2_min, self.p_2_max), (self.p_3_min, self.p_3_max),
                                   (self.beta_1_min, self.beta_1_max),(self.beta_2_min, self.beta_2_max),
                                   (self.beta_3_min, self.beta_3_max), (self.eposilon_min, self.eposilon_max),
                                   (self.theta_min, self.theta_max), (self.alpha_min, self.alpha_max),
                                   (self.eta_min, self.eta_max), (self.gamma_min, self.gamma_max),
                                   (self.Q_min, self.Q_max), (self.Lambda_min, self.Lambda_max)])
        print(optimal)
        # 结果预测
        q_1, q_2, q_3, p_1,p_2,p_3, beta_1,beta_2,beta_3, eposilon, theta, alpha, eta, gamma, Q, Lambda = optimal.x
        new_index, extended_S, extended_Sq, extended_E, extended_Ehq, extended_Ecq, extended_A, extended_I, extended_R, prediction = self.predict(
            q_1, q_2, q_3, p_1,p_2,p_3, beta_1,beta_2,beta_3, eposilon, theta, alpha, eta, gamma, Q, Lambda, S, Sq, E, Ehq, Ecq, A, I, R,
            self.s_0,
            self.sq_0, self.e_0, self.ehq_0, self.ecq_0, self.i_0, self.a_0, self.r_0)
        df = pd.DataFrame(
            {'Real S': extended_S, 'Real R': extended_R, 'Real A': extended_A, 'Real I': extended_I,
             'Converge S': prediction.y[0], 'Converge Sq': prediction.y[1], 'Converge E': prediction.y[2],
             'Converge Ehq': prediction.y[3], 'Converge Ecq': prediction.y[4], 'Converge I': prediction.y[5],
             'Converge A': prediction.y[6], 'Converge R': prediction.y[7]},
            index=new_index)
        print(f" q_1={q_1:.8f}, q_2={q_2:.8f}")
        print(f" q_3={q_3:.8f}, p_1={p_1:.8f}")
        print(f" p_2={p_2:.8f}, p_3={p_3:.8f}")
        print(f" beta_1={beta_1:.8f}, beta_2={beta_2:.8f}")
        print(f" beta_3={beta_3:.8f}, eposilon={eposilon:.8f}")
        print(f" theta={theta:.8f}, alpha={alpha:.8f}")
        print(f" eta={eta:.8f}, gamma={gamma:.8f}")
        print(f" Q={Q:.8f}, Lambda={Lambda:.8f}")
        result = dict()
        result['q_1'] = q_1
        result['q_2'] = q_2
        result['q_3'] = q_3
        result['t_1'] = self.t_1
        result['t_2'] = self.t_2
        result['p_1'] = p_1
        result['p_2'] = p_2
        result['p_3'] = p_3
        result['beta_1'] = beta_1
        result['beta_2'] = beta_2
        result['beta_3'] = beta_3
        result['eposilon'] = eposilon
        result['theta'] = theta
        result['alpha'] = alpha
        result['eta'] = eta
        result['gamma'] = gamma
        result['Q'] = Q
        result['Lambda'] = Lambda
        fw = open(self.parameter_output, 'w', encoding='utf-8')
        fw.write(f"q_1={q_1:.8f}\nq_2={q_2:.8f}\n")
        fw.write(f"q_3={q_3:.8f}\np_1={p_1:.8f}\n")
        fw.write(f"p_2={p_2:.8f}\np_3={p_3:.8f}\n")
        fw.write(f"beta_1={beta_1:.8f}\nbeta_2={beta_2:.8f}\n")
        fw.write(f"beta_3={beta_3:.8f}\neposilon={eposilon:.8f}\n")
        fw.write(f"theta={theta:.8f}\nalpha={alpha:.8f}\n")
        fw.write(f"eta={eta:.8f}\ngamma={gamma:.8f}\n")
        fw.write(f"Q={Q:.8f}\nLambda={Lambda:.8f}\n")
        fw.close()
        return df, result


def loss(point, t_1, t_2, I_loss_rate, A_loss_rate, R_loss_rate, A, I, R, s_0, sq_0, e_0, ehq_0, ecq_0, i_0, a_0, r_0):
    size = len(I)
    q_1, q_2, q_3, p_1,p_2,p_3, beta_1,beta_2,beta_3, eposilon, theta, alpha, eta, gamma, Q, Lambda = point

    def SEIR(t, y):
        '''
        return: ds,de,di,dr
        '''
        S = y[0]
        Sq = y[1]
        E = y[2]
        Ehq = y[3]
        Ecq = y[4]
        I = y[5]
        A = y[6]
        R = y[7]

        q = q_1 if t < t_1 else q_2 if t < t_2 + t_1 else q_3
        p = p_1 if t < t_1 else p_2 if t < t_2 + t_1 else p_3
        beta = beta_1 if t < t_1 else beta_2 if t < t_2 + t_1 else beta_3

        Ehq_in = (1 - q) * p * beta * (I + eposilon * E + theta * A) * S
        Ehq_out = alpha * eta * Ehq + alpha * (1 - eta) * Ehq

        E_in = (1 - q) * (1 - p) * beta * (I + eposilon * E + theta * A) * S
        E_out = alpha * eta * E + alpha * (1 - eta) * E

        Ecq_in = q * beta * (I + eposilon * E + theta * A) * S
        Ecq_out = alpha * eta * Ecq + alpha * (1 - eta) * Ecq

        S_in = Lambda * Sq
        S_out = Ehq_in + E_in + Ecq_in + Q

        Sq_out = Lambda * Sq
        Sq_in = Q

        I_in = alpha * eta * (Ehq + E + Ecq)
        A_in = alpha * (1 - eta) * (Ehq + E + Ecq)

        I_out = gamma * I
        A_out = gamma * A

        R_in = I_out + A_out

        return [S_in - S_out,
                Sq_in - Sq_out,
                E_in - E_out,
                Ehq_in - Ehq_out,
                Ecq_in - Ecq_out,
                I_in - I_out,
                A_in - A_out,
                R_in]

    solution = solve_ivp(SEIR, [0, size], [s_0, sq_0, e_0, ehq_0, ecq_0, i_0, a_0, r_0], t_eval=np.arange(0, size, 1),
                         vectorized=True)
    l1 = np.sqrt(np.mean((solution.y[5] - I) ** 2))
    l2 = np.sqrt(np.mean((solution.y[6] - A) ** 2))
    l3 = np.sqrt(np.mean((solution.y[7] - R) ** 2))
    a1 = I_loss_rate
    a2 = A_loss_rate
    a3 = R_loss_rate
    loss_item = a1 * l1 + a2 * l2 + a3 * l3
    print(loss_item)
    return loss_item
